package com.example.demo.leetcode.classics150;

import java.util.*;

/**
 * ******************************************************
 *
 * @author liugh9
 * @version 1.0
 * @classname _123查找和最小的K对数字
 * @description
 * @date 2023/08/17 20:57
 * <p>
 * ******************************************************
 */
public class _123查找和最小的K对数字 {


    /**
     * 我们可以使用优先级队列来实现，队列中保存的是 [index1, index2]，
     * 分别表示 nums1 的索引和 nums2 的索引，初始时把 [0,0]、[1, 0]、[2, 0]、…… 入队，
     * 即让 nums2 的索引全部从 0 开始，每次弹出 nums1[index1] + nums2[index2] 较小者。
     * 弹出之后，再把 index2 后移一位继续添加到优先级队列中，依次往复，最终弹出 k 次就是我们的结果。
     *
     *
     * @param nums1
     * @param nums2
     * @param k
     * @return
     */
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // nums1对nums2的索引映射，index1->index2，初始值都是0
        int[] indices2 = new int[nums1.length];

        // 优先级队列，相加之和较小的排在前面
        PriorityQueue<Integer> heap = new PriorityQueue<>(Comparator.comparingInt(i -> nums1[i] + nums2[indices2[i]]));

        // 把nums1的索引都加入到优先级队列中
        // 优化：最多加 k 个即可
        for (int i = 0; i < Math.min(nums1.length, k); i++) {
            heap.offer(i);
        }

        List<List<Integer>> ans = new ArrayList<>();

        // 最多取 k 次
        while (k-- > 0 && !heap.isEmpty()) {
            // nums1的索引
            int index1 = heap.poll();
            // nums2的索引
            int index2 = indices2[index1];
            // 添加结果
            ans.add(Arrays.asList(nums1[index1], nums2[index2]));

            // 如果nums2的索引没到头，可以继续加入队列
            if (++indices2[index1] < nums2.length) {
                heap.offer(index1);
            }
        }

        return ans;
    }


    public List<List<Integer>> kSmallestPairs1(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> res = new LinkedList<>();
        int m = nums1.length;
        int n = nums2.length;
        int size = m * n;
        List<int[]> newNums = new ArrayList<>(size);
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                newNums.add(new int[]{nums1[i], nums2[j]});
            }
        }
        PriorityQueue<int[]> pq = new PriorityQueue<>((a, b) -> (b[0] + b[1] - a[0] - a[1]));
        for (int[] newNum : newNums) {
            pq.offer(newNum);
            if (pq.size() > k) {
                pq.poll();
            }
        }
        while (!pq.isEmpty()) {
            int[] poll = pq.poll();
            res.add(new ArrayList<Integer>() {{
                add(poll[0]);
                add(poll[1]);
            }});
        }
        return res;
    }

}
